// Copyright 2022 Chen Jun
// Licensed under the MIT License.

//#include <cv_bridge/cv_bridge.h>
//#include <rmw/qos_profiles.h>
//#include <tf2/LinearMath/Matrix3x3.h>
//#include <tf2/convert.h>

//#include <ament_index_cpp/get_package_share_directory.hpp>
//#include <image_transport/image_transport.hpp>
#include <opencv2/core.hpp>
#include <opencv2/imgproc.hpp>
//#include <rclcpp/duration.hpp>
//#include <rclcpp/qos.hpp>
//#include <tf2_geometry_msgs/tf2_geometry_msgs.hpp>

// STD
#include <algorithm>
#include <map>
#include <memory>
#include <string>
#include <vector>

#include "armor.hpp"
#include "detector_node.hpp"

//#include "armor_detector/armor.hpp"
//#include "armor_detector/detector_node.hpp"

namespace rm_auto_aim
{
//ArmorDetectorNode::ArmorDetectorNode(const rclcpp::NodeOptions & options)
//: Node("armor_detector", options)
//{
ArmorDetectorNode::ArmorDetectorNode()
{
//  RCLCPP_INFO(this->get_logger(), "Starting DetectorNode!");
  std::cout << "Starting DetectorNode!" << std::endl;

  // Detector
  detector_ = initDetector();

  // Armors Publisher
//  armors_pub_ = this->create_publisher<auto_aim_interfaces::msg::Armors>(
//    "/detector/armors", rclcpp::SensorDataQoS());

  // Visualization Marker Publisher
  // See http://wiki.ros.org/rviz/DisplayTypes/Marker
//  armor_marker_.ns = "armors";
//  armor_marker_.action = visualization_msgs::msg::Marker::ADD;
//  armor_marker_.type = visualization_msgs::msg::Marker::CUBE;
//  armor_marker_.scale.x = 0.05;
//  armor_marker_.scale.z = 0.125;
//  armor_marker_.color.a = 1.0;
//  armor_marker_.color.g = 0.5;
//  armor_marker_.color.b = 1.0;
//  armor_marker_.lifetime = rclcpp::Duration::from_seconds(0.1);

//  text_marker_.ns = "classification";
//  text_marker_.action = visualization_msgs::msg::Marker::ADD;
//  text_marker_.type = visualization_msgs::msg::Marker::TEXT_VIEW_FACING;
//  text_marker_.scale.z = 0.1;
//  text_marker_.color.a = 1.0;
//  text_marker_.color.r = 1.0;
//  text_marker_.color.g = 1.0;
//  text_marker_.color.b = 1.0;
//  text_marker_.lifetime = rclcpp::Duration::from_seconds(0.1);

//  marker_pub_ =
//    this->create_publisher<visualization_msgs::msg::MarkerArray>("/detector/marker", 10);

  // Debug Publishers
//  debug_ = this->declare_parameter("debug", false);
  debug_ = true;
  if (debug_) {
    createDebugPublishers();
  }

  // Debug模式相关订阅
  // Debug param change moniter
//  debug_param_sub_ = std::make_shared<rclcpp::ParameterEventHandler>(this);
//  debug_cb_handle_ =
//    debug_param_sub_->add_parameter_callback("debug", [this](const rclcpp::Parameter & p) {
//      debug_ = p.as_bool();
//      debug_ ? createDebugPublishers() : destroyDebugPublishers();
//    });

  // 相机相关的订阅
//  cam_info_sub_ = this->create_subscription<sensor_msgs::msg::CameraInfo>(
//    "/camera_info", rclcpp::SensorDataQoS(),
//    // 回调函数，lambda表达式
//    [this](sensor_msgs::msg::CameraInfo::ConstSharedPtr camera_info) {
//      cam_center_ = cv::Point2f(camera_info->k[2], camera_info->k[5]);
//      cam_info_ = std::make_shared<sensor_msgs::msg::CameraInfo>(*camera_info);
//      pnp_solver_ = std::make_unique<PnPSolver>(camera_info->k, camera_info->d);
//      cam_info_sub_.reset(); // 在使用 reset() 方法时，我们需要保证订阅者在重置之前已经被销毁或不再使用。
//    });
// TODO 初始化 cam_center_ pnp_solver_

  // 相机拉流图片结果相关的订阅
//  img_sub_ = this->create_subscription<sensor_msgs::msg::Image>(
//    "/image_raw", rclcpp::SensorDataQoS(),
//    std::bind(&ArmorDetectorNode::imageCallback, this, std::placeholders::_1));
// TODO 对相机拉流图像使用imageCallback
}

//void ArmorDetectorNode::imageCallback(const sensor_msgs::msg::Image::ConstSharedPtr img_msg)
//{
void ArmorDetectorNode::imageCallback(const cv::Mat &src)
{
  auto armors = detectArmors(src);

  if (pnp_solver_ != nullptr) {
//    armors_msg_.header = armor_marker_.header = text_marker_.header = img_msg->header;
//    armors_msg_.armors.clear();
//    marker_array_.markers.clear();
//    armor_marker_.id = 0; // id：Marker 的唯一标识符
//    text_marker_.id = 0;

//    auto_aim_interfaces::msg::Armor armor_msg;
    RotatedArmor rotatedArmor;
    for (const auto & armor : armors) {
      cv::Mat rvec, tvec;
      bool success = pnp_solver_->solvePnP(armor, rvec, tvec);
      if (success) {
        // Fill basic info
//        armor_msg.type = ARMOR_TYPE_STR[static_cast<int>(armor.type)];
//        armor_msg.number = armor.number;
          rotatedArmor.armorType = armor.type;
          rotatedArmor.tag_id = armor.number;

        // Fill pose
//        armor_msg.pose.position.x = tvec.at<double>(0);
//        armor_msg.pose.position.y = tvec.at<double>(1);
//        armor_msg.pose.position.z = tvec.at<double>(2);
        rotatedArmor.position.x = tvec.at<double>(0);
        rotatedArmor.position.y = tvec.at<double>(1);
        rotatedArmor.position.z = tvec.at<double>(2);

        // rvec to 3x3 rotation matrix
        cv::Mat rotation_matrix;
        cv::Rodrigues(rvec, rotation_matrix);
        // rotation matrix to quaternion
//        tf2::Matrix3x3 tf2_rotation_matrix(
//          rotation_matrix.at<double>(0, 0), rotation_matrix.at<double>(0, 1),
//          rotation_matrix.at<double>(0, 2), rotation_matrix.at<double>(1, 0),
//          rotation_matrix.at<double>(1, 1), rotation_matrix.at<double>(1, 2),
//          rotation_matrix.at<double>(2, 0), rotation_matrix.at<double>(2, 1),
//          rotation_matrix.at<double>(2, 2));
//        tf2::Quaternion tf2_q;
//        tf2_rotation_matrix.getRotation(tf2_q);
//        armor_msg.pose.orientation = tf2::toMsg(tf2_q);

        // Fill the distance to image center
//        armor_msg.distance_to_image_center = pnp_solver_->calculateDistanceToCenter(armor.center);
        rotatedArmor.distance_to_image_center = pnp_solver_->calculateDistanceToCenter(armor.center);

          // Fill the markers
//        armor_marker_.id++;
//        armor_marker_.scale.y = armor.type == ArmorType::SMALL ? 0.135 : 0.23;
//        armor_marker_.pose = armor_msg.pose;
//        text_marker_.id++;
//        text_marker_.pose.position = armor_msg.pose.position;
//        text_marker_.pose.position.y -= 0.1;
//        text_marker_.text = armor.classfication_result;
//        armors_msg_.armors.emplace_back(armor_msg);
//        marker_array_.markers.emplace_back(armor_marker_);
//        marker_array_.markers.emplace_back(text_marker_);
//TODO 建立识别到的所有装甲板容器
      } else {
//        RCLCPP_WARN(this->get_logger(), "PnP failed!");
        std::cout << "PnP failed!" << std::endl;
      }
    }

    // Publishing detected armors
//    armors_pub_->publish(armors_msg_);
    // TODO 传出识别到的装甲板

    // Publishing marker
    publishMarkers();
  }
}

std::unique_ptr<Detector> ArmorDetectorNode::initDetector()
{
    // param_desc 描述参数的类型、范围
//  rcl_interfaces::msg::ParameterDescriptor param_desc;
//  param_desc.integer_range.resize(1); 数组类型
//  param_desc.integer_range[0].step = 1;
//  param_desc.integer_range[0].from_value = 0;
//  param_desc.integer_range[0].to_value = 255;
//  int binary_thres = declare_parameter("binary_thres", 160, param_desc);
  binary_thres = 220;

//  param_desc.description = "0-RED, 1-BLUE";
//  param_desc.integer_range[0].from_value = 0;
//  param_desc.integer_range[0].to_value = 1;
//  auto detect_color = declare_parameter("detect_color", RED, param_desc);
  detect_color = RED;
    std::cout << "light.color = " << detect_color << std::endl;
//  Detector::LightParams l_params = {
//    .min_ratio = declare_parameter("light.min_ratio", 0.1),
//    .max_ratio = declare_parameter("light.max_ratio", 0.4),
//    .max_angle = declare_parameter("light.max_angle", 40.0)};
  Detector::LightParams l_params = {
    .min_ratio = 0.1,
    .max_ratio = 0.4,
    .max_angle = 40.0};

//  Detector::ArmorParams a_params = {
//    .min_light_ratio = declare_parameter("armor.min_light_ratio", 0.7),
//    .min_small_center_distance = declare_parameter("armor.min_small_center_distance", 0.8),
//    .max_small_center_distance = declare_parameter("armor.max_small_center_distance", 3.2),
//    .min_large_center_distance = declare_parameter("armor.min_large_center_distance", 3.2),
//    .max_large_center_distance = declare_parameter("armor.max_large_center_distance", 5.5),
//    .max_angle = declare_parameter("armor.max_angle", 35.0)};
  Detector::ArmorParams a_params = {
    .min_light_ratio = 0.7,
    .min_small_center_distance = 0.8,
    .max_small_center_distance = 3.2,
    .min_large_center_distance = 3.2,
    .max_large_center_distance = 5.5,
    .max_angle = 35.0};

  auto detector = std::make_unique<Detector>(binary_thres, detect_color, l_params, a_params);

  // Init classifier
//  auto pkg_path = ament_index_cpp::get_package_share_directory("armor_detector");
  std::string pkg_path = ".."; // 数字分类的模型文件
  auto model_path = pkg_path + "/model/mlp.onnx";
  auto label_path = pkg_path + "/model/label.txt";
//  double threshold = this->declare_parameter("classifier_threshold", 0.7);
  classifier_threshold = 0.8; // classifier_threshold
//  std::vector<std::string> ignore_classes =
//    this->declare_parameter("ignore_classes", std::vector<std::string>{"negative"});
  std::vector<ArmorID> ignore_classes = std::vector<ArmorID>{ArmorID::NEGATIVE}; //ignore_classes
  detector->classifier =
//    std::make_unique<NumberClassifier>(model_path, label_path, classifier_threshold, ignore_classes);
    std::make_unique<NumberClassifier>(model_path, classifier_threshold, ignore_classes);

  return detector;
}

//std::vector<Armor> ArmorDetectorNode::detectArmors(
//  const sensor_msgs::msg::Image::ConstSharedPtr & img_msg)
//{

std::vector<Armor> ArmorDetectorNode::detectArmors(const cv::Mat &src)
{
    // img类型是cv::Mat
  // Convert ROS img to cv::Mat
//  auto img = cv_bridge::toCvShare(img_msg, "rgb8")->image;
    auto img = src.clone();

  // Update params
  detector_->binary_thres = binary_thres;
  detector_->detect_color = detect_color;
  detector_->classifier->threshold = classifier_threshold;

  auto armors = detector_->detect(img);
//std::cout << "detect \n";
//  auto final_time = this->now();
//  auto latency = (final_time - img_msg->header.stamp).seconds() * 1000;
//  RCLCPP_DEBUG_STREAM(this->get_logger(), "Latency: " << latency << "ms");
//    std::cout << "Latency: " << latency << "ms" << std::endl;
    // TODO 计算识别时间
    // 当前时间 减去 拉流图像的时间

  // Publish debug info
  if (debug_) {
//    binary_img_pub_.publish(
//      cv_bridge::CvImage(img_msg->header, "mono8", detector_->binary_img).toImageMsg());
      cv::imshow("binary_img", detector_->binary_img);
// TODO 显示二值化图像
    // Sort lights and armors data by x coordinate
//    std::sort(
//      detector_->debug_lights.data.begin(), detector_->debug_lights.data.end(),
//      [](const auto & l1, const auto & l2) { return l1.center_x < l2.center_x; });
//    std::sort(
//      detector_->debug_armors.data.begin(), detector_->debug_armors.data.end(),
//      [](const auto & a1, const auto & a2) { return a1.center_x < a2.center_x; });
//
//    lights_data_pub_->publish(detector_->debug_lights);
//    armors_data_pub_->publish(detector_->debug_armors);

    if (!armors.empty()) {
      auto all_num_img = detector_->getAllNumbersImage();
      // TODO 显示识别到的数字
//      number_img_pub_.publish(
//        *cv_bridge::CvImage(img_msg->header, "mono8", all_num_img).toImageMsg());
        cv::imshow("all_num_img", all_num_img);
    }

    detector_->drawResults(img);
    // Draw camera center
    cv::circle(img, cam_center_, 5, cv::Scalar(255, 0, 0), 2);
    // Draw latency
    // TODO 绘制识别时间
//    std::stringstream latency_ss;
//    latency_ss << "Latency: " << std::fixed << std::setprecision(2) << latency << "ms";
//    auto latency_s = latency_ss.str();
//    cv::putText(
//      img, latency_s, cv::Point(10, 30), cv::FONT_HERSHEY_SIMPLEX, 1.0, cv::Scalar(0, 255, 0), 2);

//    result_img_pub_.publish(cv_bridge::CvImage(img_msg->header, "rgb8", img).toImageMsg());
//TODO 显示识别后的结果图像
cv::imshow("result_img", img);
  }

  return armors;
}

// 创建debug话题发布者
void ArmorDetectorNode::createDebugPublishers()
{
//  lights_data_pub_ =
//    this->create_publisher<auto_aim_interfaces::msg::DebugLights>("/detector/debug_lights", 10);
//  armors_data_pub_ =
//    this->create_publisher<auto_aim_interfaces::msg::DebugArmors>("/detector/debug_armors", 10);
//
//  binary_img_pub_ = image_transport::create_publisher(this, "/detector/binary_img");
//  number_img_pub_ = image_transport::create_publisher(this, "/detector/number_img");
//  result_img_pub_ = image_transport::create_publisher(this, "/detector/result_img");
}

// 销毁debug话题发布者
void ArmorDetectorNode::destroyDebugPublishers()
{
//  lights_data_pub_.reset();
//  armors_data_pub_.reset();
//
//  binary_img_pub_.shutdown();
//  number_img_pub_.shutdown();
//  result_img_pub_.shutdown();
}

void ArmorDetectorNode::publishMarkers()
{
//  using Marker = visualization_msgs::msg::Marker;
//  armor_marker_.action = armors_msg_.armors.empty() ? Marker::DELETE : Marker::ADD;
//  marker_array_.markers.emplace_back(armor_marker_);
//  marker_pub_->publish(marker_array_);
}

}  // namespace rm_auto_aim

//#include "rclcpp_components/register_node_macro.hpp"

// Register the component with class_loader.
// This acts as a sort of entry point, allowing the component to be discoverable when its library
// is being loaded into a running process.
//RCLCPP_COMPONENTS_REGISTER_NODE(rm_auto_aim::ArmorDetectorNode)
